import numpy as np
import torch


def check_numpy_to_torch(x):
    if isinstance(x, np.ndarray):
        return torch.from_numpy(x).float(), True
    return x, False


def rotate_points(points, R):
    """
    Args:
        points: (B*N1, N2, 3 + C)
        R: (B*N1, 3, 3)
    Returns:
    """

    points, is_numpy = check_numpy_to_torch(points)
    R, _ = check_numpy_to_torch(R)

    points_rot = torch.matmul(points[:, :, 0:3], R.transpose(-1, -2).contiguous())
    points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
    return points_rot.numpy() if is_numpy else points_rot

